from torchvision import transforms
import numpy as np
import torch

class HilightCrop:
    def __init__(self, width):
        self.width = width

    def __call__(self, pic):
        col_max = np.max(pic, 0)
        maxmax = sorted([(x, i) for i, x in enumerate(col_max)])
        crop = [-1, -1]
        crop_size = -1
        for x, i in reversed(maxmax):
            if crop_size == -1:
                crop[0] = i
                crop[1] = i
            else:
                if i < crop[0]:
                    crop[0] = i
                elif i > crop[1]:
                    crop[1] = i
            crop_size = crop[1] - crop[0] + 1
            if crop_size >= self.width:
                break

        if crop_size > self.width:
            margin = crop_size - self.width
            a, res = divmod(margin, 2)
            crop[0] += a
            crop[1] -= a
            crop[1] -= res
        if crop[0] > crop[1]:
            crop[0] = crop[1]

        return pic.crop((crop[0], 0, crop[1] + 1, pic.size[1]))

transform = transforms.Compose([
    HilightCrop(240),
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Resize((96, 96), antialias=True),  # 调整图像大小
    transforms.Normalize((0.5,), (0.5,))  # 归一化
])